#!/bin/python3
# -*- coding: UTF-8
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.colors
import matplotlib.dates as mdates
from matplotlib.ticker import ScalarFormatter
from matplotlib.ticker import LogFormatter
from matplotlib.ticker import MultipleLocator
from matplotlib.ticker import FormatStrFormatter
from matplotlib.font_manager import findfont, FontProperties  
import matplotlib.font_manager as font_manager
import matplotlib.ticker as mtick

import seaborn as sns

sns.set_theme(style="ticks")
sns.set_context("paper", font_scale=1.5)
plt.rcParams['xtick.direction'] = 'in'#将x的刻度线方向设置向内
plt.rcParams['ytick.direction'] = 'in'#将y的刻度方向设置向内
plt.rcParams['font.family']=['Sarasa Mono SC'] #用来正常显示中文标签
plt.rcParams['axes.xmargin'] = 0
plt.rcParams['axes.ymargin'] = 0

txt_workloads = "测试负载"
txt_executiontime = "执行时间(按照wfslab归一化)"
txt_rss = "最大RSS(按照wfslab归一化)"


bench_list1 = ['cfrac', 'espresso', 'barnes', 'redis', 'leanN', 'z3']
bench_list2 = ['rbstress1', 'rbstressN', 'gs', 'lua', 'cxqueue', 'rocksdb']
bench_list3 = ['larsonN', 'alloc-test1', 'alloc-testN', 'sh6benchN', 'sh8benchN', 'xmalloc-testN']
bench_list3_aarch64 = ['larsonN', 'alloc-test1', 'alloc-testN', 'sh6benchN', 'xmalloc-testN']	# sh8跑不了
bench_list4 = ['cache-scratch1', 'cache-scratchN', 'rptestN', 'mstressN', 'glibc-simple', 'glibc-thread']


bencn_baseline1 = 'wfslab'
# bencn_baseline2 = 'ccsync'
hatch_list = [''] 

# TODO: 色表修改
color_base_max = 1.0


colortbls = [	
		(color_base_max * 13.0/14.0,  color_base_max* 13.0/14.0, color_base_max*13.0/14.0), # sys
		(color_base_max * 12.0/14.0,  color_base_max* 12.0/14.0, color_base_max*12.0/14.0), # hd
		(color_base_max * 11.0/14.0,  color_base_max* 11.0/14.0, color_base_max*11.0/14.0), # tc
		(color_base_max * 10.0/14.0,  color_base_max* 10.0/14.0, color_base_max*10.0/14.0), # je
		(color_base_max * 9.0/14.0,  color_base_max* 9.0/14.0, color_base_max*9.0/14.0), # lp
		(color_base_max * 8.0/14.0,  color_base_max* 8.0/14.0, color_base_max*8.0/14.0), # tbb
		(color_base_max * 147/255,  color_base_max* 207/255, color_base_max*201/255), # rp 无锁
		(color_base_max * 7.0/14.0,  color_base_max* 7.0/14.0, color_base_max*7.0/14.0), # mesh 锁
		(color_base_max * 147/255,  color_base_max* 207/255, color_base_max*201/255), # mi 无锁
		(color_base_max * 147/255,  color_base_max* 207/255, color_base_max*201/255), # sn 无锁
		(color_base_max * 250/255,  color_base_max* 127/255, color_base_max*111/255), # wfspan
		(color_base_max * 250/255,  color_base_max* 127/255, color_base_max*111/255), # wfspan_wfqp
		(color_base_max * 130/255,  color_base_max* 176/255, color_base_max*210/255), # tlsf
		(color_base_max * 130/255,  color_base_max* 176/255, color_base_max*210/255), # tlsf_sffwd
		# (color_base_max * 231/255,  color_base_max* 218/255, color_base_max*210/255),
		(color_base_max * 190/255,  color_base_max* 184/255, color_base_max*220/255), # hslab
		(color_base_max * 250/255,  color_base_max* 127/255, color_base_max*111/255), # wfslab
		]
		
hatchtbls = ['']
		
# colortbls=[(0/255, 45/255, 106/255),(0x99/255, 0xc5/255,0x56/255), (0/255, 92/255, 171/255), (220/255, 176/255, 243/255), (255/255, 195/255, 37/255), (0xfd/255, 0xaf/255, 0xab/255),  (227/255, 67/255, 35/255),
# (0/255, 147/255, 43/255),(0x99/255, 0xa8/255,0x56/255), (0/255, 92/255, 121/255), (227/255, 27/255, 35/255)
# ]
cl_blk = (7/255,7/255,7/255,1)
time_column = 0
rss_column = 1

# y轴上限
avg_ylimit = 1
size_ylimit = 0.5
size_ybot = -0.5

size_ylimitrss = 0.5
size_ybotrss = -0.5

accuracy = 1

keynote_font = 9

#  搞一个数组,这里画了就+1
def autolabel(ax,rects, ylimit, rect_table):
	"""Attach a text label above each bar in *rects*, displaying its height."""
	t = 0
	for rect in rects:
		# print (rect)
		height = rect.get_height()
		#print(rect.get_height())
		top_height = 0.1
		# if (rect.get_height()>=ylimit):
		# 	# print(0 - (top_height * rect_table[t]))
		# 	ax.annotate('{:3.2f}'.format(1+height),
		# 				xy=(rect.get_x() + rect.get_width() / 2, size_ylimit),
		# 				xytext=(0, 1),  # 3 points vertical offset
		# 				fontsize=keynote_font
		# 				,textcoords="offset points",
		# 				rotation = 90,
		# 				ha='center', va='bottom')
		# 	# ax.vlines(x=rect.get_x() + rect.get_width() / 2, ymin= size_ybot + (top_height * (rect_table[t] + 1)), ymax=0, alpha=0.618, linewidth=0.5, linestyles='dashed', color = cl_blk)
		# 	rect_table[t] += 1
		# if (rect.get_height()>=0) :
		# 	ax.annotate('{:3.2f}'.format(1+height),
		# 				xy=(rect.get_x() + rect.get_width() / 2, 0),
		# 				xytext=(0, -21),  # 3 points vertical offset
		# 				fontsize=keynote_font
		# 				,textcoords="offset points",
		# 				rotation = 90,
		# 				ha='center', va='bottom')
		# else : 
		# 	ax.annotate('{:3.2f}'.format(1+height),
		# 				xy=(rect.get_x() + rect.get_width() / 2, size_ylimit),
		# 				xytext=(0, 1),  # 3 points vertical offset
		# 				fontsize=keynote_font
		# 				,textcoords="offset points",
		# 				rotation = 90,
		# 				ha='center', va='bottom')
		# 	# ax.vlines(x=rect.get_x() + rect.get_width() / 2, ymin= size_ybot + (top_height * (rect_table[t] + 1)), ymax=0, alpha=0.618, linewidth=0.5, linestyles='dashed', color = cl_blk)
		# 	rect_table[t] += 1
		# 	# print(rect_table)
		# ax.annotate('{:3.2f}'.format(1+height),
		# 				xy=(rect.get_x() + rect.get_width() / 2, ylimit),
		# 				xytext=(0, 1),  # 3 points vertical offset
		# 				fontsize=keynote_font
		# 				,textcoords="offset points",
		# 				rotation = 90,
		# 				ha='center', va='bottom')
		if (rect.get_height()>=ylimit):
			ax.annotate("%6.1fx" % (1+height),
						# xy=(rect.get_x() + rect.get_width() / 2, size_ylimit),
						xy=(rect.get_x() + rect.get_width() / 2, 0),
						xytext=(0, 1),  # 3 points vertical offset
						font = "Noto Sans",
						fontsize = keynote_font - 3,
						rotation = 90,
						textcoords="offset points",
						ha='center', va='bottom')
		t+=1

# 读文件
def read_data(file_path):
	data = np.loadtxt(file_path,dtype=str)
	# print (data)
	return data

def get_bench_data(data, bench_name):
	ret = data[data[:, 0] == bench_name]
	# 裁减
	return ret

def __get_bench_data_x_y(data, bench_name):
	ret = get_bench_data(data, bench_name)
	x_name = ret[:, 1]
	y = ret[:, 2:]
	# print (y)
	return x_name, y

def get_bench_data_x_y(data, bench_name, column, typ):
	x_name, y = __get_bench_data_x_y(data, bench_name)
	x_num = np.unique(x_name).size
	x_name_row = int(x_name.size / x_num)
	x_name = x_name.reshape(x_name_row, x_num)

	# print (x_name)
	# print (y)
	y = y[:, column].astype(typ)
	y = np.array(y)

	y = y.reshape(x_name_row, x_num)
	# y.
	return x_name, np.mean(y, axis=0), np.std(y, axis=0)


def plot_subfig_avg(data_x, data_avg_y, data_std, axs, ylimit, benchline1):
	x_length = len(data_x)
			   # the label locations

	width = 0.9

	rel = np.where(data_x == benchline1)		 # use bencn_baseline
	rel = rel[0][0]

	# print (data_x[rel])
	# print (data_avg_y[0][rel])
	# print (np.true_divide(data_avg_y[0], data_avg_y[0][rel]))

	rects = []
	column = len(data_avg_y[:])
	x = np.arange(column)   
	

	div = 1 / (x_length + 1)
	medium = 1 / 2

	# print(data_x)
	# true_x = - width * (medium - (x + 1) * div)
	# print (true_x)
	# x[i] - width*(medium - (i + 1) * div)
	# print ( column)
	# print (column)
	# # print (data_x)
	# print (np.true_divide(data_avg_y[:,0], data_avg_y[:, rel]))
	color_cnt = 0
	for i in range(0, x_length):
		# axs.vlines(x=x-0.5 , ymin= -1, ymax=ylimit, linewidth=0.5, linestyles='dashed', alpha=0.618)
		# axs.vlines(x=x+0.5 , ymin= -1, ymax=ylimit, linewidth=0.5, linestyles='dashed',  alpha=0.618)
		# print (data_avg_y[i])
		# print (data_avg_y[:,i])
		# print ( np.true_divide(data_avg_y[:,i], data_avg_y[:, rel]))
		# 这是横着画的
		if data_x[i] in hatch_list:
			hatchval = 'xx'
			color = colortbls[(color_cnt - 1)%len(colortbls)]
		else:
			hatchval = hatchtbls[color_cnt%len(hatchtbls)]
			color = colortbls[color_cnt%len(colortbls)]
			color_cnt += 1
		
		print (data_x[i])
		print (np.true_divide(data_avg_y[:,i], data_avg_y[:, rel]) - 1)
		# height = np.true_divide(data_avg_y[:,i], data_avg_y[:, rel]) - 1
		rects.append(axs.bar( x - width*(medium - (i + 1) * div), 
				np.true_divide(data_avg_y[:,i], data_avg_y[:, rel]) - 1, 
				width * div, 
					yerr = np.true_divide(data_std[:,i], data_avg_y[:, rel]), 
						label=data_x[i],
							color=color, 
								alpha=0.9, capsize=1, error_kw={"elinewidth" : 1, "alpha" : 0.8}, hatch=hatchval))
	return rects


def set_axs_label(axs, data_x, data_avg_y, rects, xlabel, ylabel, ybotlimit, ytoplimit, benchlist):
	x_length = len(data_x)
	column = len(data_avg_y)
	x = np.arange(column)
	
	axs.set_xlabel(xlabel, fontsize=keynote_font)
	axs.set_ylabel(ylabel, fontsize=keynote_font)
	#axs.set_title('Virtual Memory Claimed Size Ratio (vs Glibc) For X86')
	axs.set_xticks(x)
	axs.set_xticklabels(benchlist, rotation=5)
	axs.yaxis.set_major_formatter(mtick.PercentFormatter(1.0))

	# 画横线区分
	idx = 0
	for xname in x:
		axs.hlines(y=0, xmin=x[idx]-0.45, xmax=x[idx]+0.45, linewidth=0.5, color = cl_blk)
		idx += 1
	# axs.hlines(y=0, xmin= x[0] - 0.5, xmax=x[len(x)-1]+ 0.5, linewidth=0.5, linestyles='dashdot')

	# axs.set_yscale("log")
	axs.set_ylim(ybotlimit ,ytoplimit)
	# y_major_locator=MultipleLocator(1)
	# axs.yaxis.set_major_locator(y_major_locator)
	# axs.yaxis.set_major_formatter(LogFormatter())
	rect_table = np.zeros(x_length)
	t=0
	for rect in rects:
		autolabel(axs, rect, ytoplimit, rect_table)
		t = t + 1
		# i+=1


def plot_figa(data_x, data_avg_y, data_std, xlabel, ylabel, ybotlimit, ytoplimit, benchlist, benchline1, name):
	# fig, axs = plt.subplots(nrows=1, ncols=2, figsize=(12, 4))
	# # plt.xticks(rotation=120)
	# rects = plot_subfig_avg(data_x, data_avg_y, data_std, axs[0], ytoplimit, benchline1)
	# set_axs_label(axs[0], data_x, data_avg_y, rects, xlabel, ylabel, ybotlimit, ytoplimit, benchlist)
	# rects = plot_subfig_avg(data_x, data_avg_y_arm, data_std_arm, axs[1], ytoplimit, benchline1)
	# set_axs_label(axs[1], data_x, data_avg_y, rects, xlabel, ylabel, ybotlimit, ytoplimit, benchlist_arm)

	# 每张一个
	fig, axs = plt.subplots(nrows=1, ncols=1, figsize=(10, 4))
	# plt.xticks(rotation=120)
	rects = plot_subfig_avg(data_x, data_avg_y, data_std, axs, ytoplimit, benchline1)
	set_axs_label(axs, data_x, data_avg_y, rects, xlabel, ylabel, ybotlimit, ytoplimit, benchlist)

	plt.legend(loc='upper center', bbox_to_anchor=(0.5, -0.1 ), frameon=False, ncol= len(data_x)/2)

	matplotlib.pyplot.subplots_adjust(left=0.06, right = 0.99, bottom = 0.2, top = 0.9, wspace = 0.10 )
	
	plt.savefig("./fig/"+name, format = "pdf")

	return fig, axs

def plot_figs(data_x, data_avg_y, data_std, data_x1, data_avg_y1, data_std1, xlabel, ylabel, ybotlimit, ytoplimit, benchlist, benchlist1, benchline1, name):
	# fig, axs = plt.subplots(nrows=1, ncols=2, figsize=(12, 4))
	# # plt.xticks(rotation=120)
	# rects = plot_subfig_avg(data_x, data_avg_y, data_std, axs[0], ytoplimit, benchline1)
	# set_axs_label(axs[0], data_x, data_avg_y, rects, xlabel, ylabel, ybotlimit, ytoplimit, benchlist)
	# rects = plot_subfig_avg(data_x, data_avg_y_arm, data_std_arm, axs[1], ytoplimit, benchline1)
	# set_axs_label(axs[1], data_x, data_avg_y, rects, xlabel, ylabel, ybotlimit, ytoplimit, benchlist_arm)

	# 每张一个
	fig, axs = plt.subplots(nrows=2, ncols=1, figsize=(12, 6))
	# plt.xticks(rotation=120)
	rects = plot_subfig_avg(data_x, data_avg_y, data_std, axs[0], ytoplimit, benchline1)
	set_axs_label(axs[0], data_x, data_avg_y, rects, "", "", ybotlimit, ytoplimit, benchlist)

	rects = plot_subfig_avg(data_x1, data_avg_y1, data_std1, axs[1], ytoplimit, benchline1)
	set_axs_label(axs[1], data_x1, data_avg_y1, rects, "", "", ybotlimit, ytoplimit, benchlist1)

	# fig.supxlabel(xlabel)
	fig.supylabel(ylabel, fontsize=keynote_font * 1.5)
	plt.legend(loc='upper center', bbox_to_anchor=(0.5, -0.06), frameon=False, ncol= len(data_x)/2)

	matplotlib.pyplot.subplots_adjust(left=0.1, right = 0.99, bottom = 0.12, top = 0.93, hspace = 0.3)
	
	plt.savefig("./fig/"+name, format = "pdf")

	return fig, axs


def plot_figs4(data_x, data_avg_y, data_std, data_x1, data_avg_y1, data_std1, data_x2, data_avg_y2, data_std2, data_x3, data_avg_y3, data_std3, xlabel, ylabel, ybotlimit, ytoplimit, benchlist, benchlist1, benchlist2, benchlist3, benchline1, name):
	# 每张一个
	fig, axs = plt.subplots(nrows=4, ncols=1, figsize=(9, 9))
	# plt.xticks(rotation=120)
	rects = plot_subfig_avg(data_x, data_avg_y, data_std, axs[0], ytoplimit, benchline1)
	set_axs_label(axs[0], data_x, data_avg_y, rects, "", "", ybotlimit, ytoplimit, benchlist)

	rects = plot_subfig_avg(data_x1, data_avg_y1, data_std1, axs[1], ytoplimit, benchline1)
	set_axs_label(axs[1], data_x1, data_avg_y1, rects, "", "", ybotlimit, ytoplimit, benchlist1)
	
	rects = plot_subfig_avg(data_x2, data_avg_y2, data_std2, axs[2], ytoplimit, benchline1)
	set_axs_label(axs[2], data_x2, data_avg_y2, rects, "", "", ybotlimit, ytoplimit, benchlist2)
	
	rects = plot_subfig_avg(data_x3, data_avg_y3, data_std3, axs[3], ytoplimit, benchline1)
	set_axs_label(axs[3], data_x3, data_avg_y3, rects, "", "", ybotlimit, ytoplimit, benchlist3)

	# fig.supxlabel(xlabel)
	fig.supylabel(ylabel, fontsize=keynote_font * 1.5, x = 0.005)
	plt.legend(loc='upper center', bbox_to_anchor=(0.45, -0.09), frameon=False, labelspacing = 0.1, handletextpad=0.2, columnspacing = 1, ncol= len(data_x)/2)

	matplotlib.pyplot.subplots_adjust(left=0.075, right = 0.995, bottom = 0.08, top = 0.99, hspace = 0.2 )
	
	plt.savefig("./fig/"+name, format = "pdf")

	return fig, axs

def calculate_data(data, bench_list, column):
	ry = []
	ry_std = []

	for bench in bench_list:
		tx, ty, tstd = get_bench_data_x_y(data, bench, column, float)
		# print (tx, ty, tstd)
		ry.append(ty)
		ry_std.append(tstd)

	rx = tx[0]
	ry = np.array(ry)
	ry_std = np.array(ry_std)
	return rx, ry, ry_std

def plot_all(data, bench_list1, bench_list2, bench_list3, bench_list4, name):
	# dx, dy, dy_std = calculate_data(data, bench_list1, 0)
	# plot_figa(dx, dy, dy_std, txt_workloads, txt_executiontime, size_ybot, size_ylimit, bench_list1, bencn_baseline1, name +"_mb_1.pdf")
	# dx, dy, dy_std = calculate_data(data, bench_list2, 0)
	# plot_figa(dx, dy, dy_std, txt_workloads, txt_executiontime, size_ybot, size_ylimit, bench_list2, bencn_baseline1, name +"_mb_2.pdf")
	# dx, dy, dy_std = calculate_data(data, bench_list3, 0)
	# plot_figa(dx, dy, dy_std, txt_workloads, txt_executiontime, size_ybot, size_ylimit, bench_list3, bencn_baseline1, name+"_mb_3.pdf")
	# dx, dy, dy_std = calculate_data(data, bench_list4, 0)
	# plot_figa(dx, dy, dy_std, txt_workloads, txt_executiontime, size_ybot, size_ylimit, bench_list4, bencn_baseline1, name+"_mb_4.pdf")

	# dx, dy, dy_std = calculate_data(data, bench_list1, 1)
	# plot_figa(dx, dy, dy_std, txt_workloads, txt_rss, size_ybotrss, size_ylimitrss, bench_list1, bencn_baseline1, name+"_mbrss_1.pdf")
	# dx, dy, dy_std = calculate_data(data, bench_list2, 1)
	# plot_figa(dx, dy, dy_std, txt_workloads, txt_rss, size_ybotrss, size_ylimitrss, bench_list2, bencn_baseline1, name+"_mbrss_2.pdf")
	# dx, dy, dy_std = calculate_data(data, bench_list3, 1)
	# plot_figa(dx, dy, dy_std, txt_workloads, txt_rss, size_ybotrss, size_ylimitrss, bench_list3, bencn_baseline1, name+"_mbrss_3.pdf")
	# dx, dy, dy_std = calculate_data(data, bench_list4, 1)
	# plot_figa(dx, dy, dy_std, txt_workloads, txt_rss, size_ybotrss, size_ylimitrss, bench_list4, bencn_baseline1, name+"_mbrss_4.pdf")
	# dx, dy, dy_std = calculate_data(data, bench_list1, 0)
	# dx1, dy1, dy_std1 = calculate_data(data, bench_list2, 0)
	# plot_figs(dx, dy, dy_std, dx1, dy1, dy_std1, txt_workloads, txt_executiontime, size_ybot, size_ylimit, bench_list1, bench_list2, bencn_baseline1, name + "_mb_1.pdf")

	# dx, dy, dy_std = calculate_data(data, bench_list3, 0)
	# dx1, dy1, dy_std1 = calculate_data(data, bench_list4, 0)
	# plot_figs(dx, dy, dy_std, dx1, dy1, dy_std1, txt_workloads, txt_executiontime, size_ybot, size_ylimit, bench_list3, bench_list4, bencn_baseline1, name + "_mb_2.pdf")

	# dx, dy, dy_std = calculate_data(data, bench_list1, 1)
	# dx1, dy1, dy_std1 = calculate_data(data, bench_list2, 1)
	# plot_figs(dx, dy, dy_std, dx1, dy1, dy_std1, txt_workloads, txt_rss, size_ybotrss, size_ylimitrss, bench_list1, bench_list2, bencn_baseline1, name + "_mbrss_1.pdf")

	# dx, dy, dy_std = calculate_data(data, bench_list3, 1)
	# dx1, dy1, dy_std1 = calculate_data(data, bench_list4, 1)
	# plot_figs(dx, dy, dy_std, dx1, dy1, dy_std1, txt_workloads, txt_rss, size_ybotrss, size_ylimitrss, bench_list3, bench_list4, bencn_baseline1, name + "_mbrss_2.pdf")
	dx, dy, dy_std = calculate_data(data, bench_list1, 0)
	dx1, dy1, dy_std1 = calculate_data(data, bench_list2, 0)
	dx2, dy2, dy_std2 = calculate_data(data, bench_list3, 0)
	dx3, dy3, dy_std3 = calculate_data(data, bench_list4, 0)
	plot_figs4(dx, dy, dy_std, dx1, dy1, dy_std1, dx2, dy2, dy_std2, dx3, dy3, dy_std3, txt_workloads, txt_executiontime, size_ybot, size_ylimit, bench_list1, bench_list2, bench_list3, bench_list4, bencn_baseline1, name + "_mb.pdf")

	dx, dy, dy_std = calculate_data(data, bench_list1, 1)
	dx1, dy1, dy_std1 = calculate_data(data, bench_list2, 1)
	dx2, dy2, dy_std2 = calculate_data(data, bench_list3, 1)
	dx3, dy3, dy_std3 = calculate_data(data, bench_list4, 1)
	plot_figs4(dx, dy, dy_std, dx1, dy1, dy_std1, dx2, dy2, dy_std2, dx3, dy3, dy_std3, txt_workloads, txt_rss, size_ybot, size_ylimit, bench_list1, bench_list2, bench_list3, bench_list4, bencn_baseline1, name + "_mbrss.pdf")

data = read_data('./mimalloc-bench_ryzen3700x.tsv')
plot_all(data, bench_list1, bench_list2, bench_list3, bench_list4, "ryzen3700x")
# 13700k
data = read_data('./mimalloc-bench_core13700k.tsv')
plot_all(data, bench_list1, bench_list2, bench_list3, bench_list4, "core13700k")
# sdm888
data = read_data('./mimalloc-bench_sdm888p.tsv')
plot_all(data, bench_list1, bench_list2, bench_list3_aarch64, bench_list4, "sdm888p")


plt.show()